package com.peng.leetcode.tree;

/**
 * PruneTree
 *
 * 814. 二叉树剪枝
 *
 * @author: lupeng6
 * @create: 2021/2/5 11:26
 */
public class PruneTree {

    public static void main(String[] args) {

    }

    public TreeNode pruneTree(TreeNode root) {
        if (root == null) {
            return null;
        }
        dfs(root);
        if (root.val == 0 && root.left == null && root.right == null) {
            return null;
        }
        return root;
    }

    void dfs(TreeNode node) {
        if (node == null) {
            return;
        }
        if (node.left != null) {
            boolean contains = isContains(node.left);
            if (!contains) {
                node.left = null;
            } else {
                dfs(node.left);
            }
        }
        if (node.right != null) {
            boolean contains = isContains(node.right);
            if (!contains) {
                node.right = null;
            } else {
                dfs(node.right);;
            }
        }
    }

    boolean isContains(TreeNode node) {
        if (node.val == 1) {
            return true;
        }
        if (node.left != null) {
            boolean contains = isContains(node.left);
            if (contains) {
                return true;
            }
        }
        if (node.right != null) {
            boolean contains = isContains(node.right);
            if (contains) {
                return true;
            }
        }
        return false;
    }
}
